package com.zzh.partnersys.ai.service;

import com.zzh.partnersys.ai.entity.dto.DocumentParseTaskStatusDTO;
import com.zzh.partnersys.ai.entity.enums.DocumentParseTaskStatus;
import com.zzh.partnersys.ai.util.RedisCacheUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import jakarta.annotation.Resource;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.concurrent.TimeUnit;

/**
 * 文档解析任务状态管理服务
 *
 * @author zzh
 */
@Slf4j
@Service
public class DocumentParseTaskService {

    @Resource
    private RedisCacheUtil redisCacheUtil;

    /**
     * Redis key 前缀
     */
    private static final String TASK_STATUS_KEY_PREFIX = "document:parse:task:";

    /**
     * 任务状态过期时间（小时）
     */
    private static final int TASK_STATUS_EXPIRE_HOURS = 24;

    /**
     * 创建任务状态
     *
     * @param taskId     任务ID
     * @param totalCount 总文档数
     */
    public void createTask(String taskId, Integer totalCount) {
        DocumentParseTaskStatusDTO statusDTO = new DocumentParseTaskStatusDTO();
        statusDTO.setTaskId(taskId);
        statusDTO.setStatus(DocumentParseTaskStatus.PROCESSING);
        statusDTO.setMessage("任务已创建，开始处理...");
        statusDTO.setTotalCount(totalCount);
        statusDTO.setCurrentIndex(0);
        statusDTO.setProgress(0);
        statusDTO.setCompletedDocs(new ArrayList<>());
        statusDTO.setFailedDocs(new ArrayList<>());
        statusDTO.setCreateTime(LocalDateTime.now());

        String key = getTaskKey(taskId);
        redisCacheUtil.setCacheObject(key, statusDTO, TASK_STATUS_EXPIRE_HOURS, TimeUnit.HOURS);
        log.info("创建文档解析任务，taskId: {}, totalCount: {}", taskId, totalCount);
    }

    /**
     * 更新任务进度
     *
     * @param taskId        任务ID
     * @param currentIndex  当前文档索引
     * @param currentDocName 当前文档名称
     * @param progress      进度（0-100）
     * @param message       状态消息
     */
    public void updateProgress(String taskId, Integer currentIndex, String currentDocName, 
                              Integer progress, String message) {
        DocumentParseTaskStatusDTO statusDTO = getTaskStatus(taskId);
        if (statusDTO == null) {
            log.warn("任务状态不存在，taskId: {}", taskId);
            return;
        }

        statusDTO.setCurrentIndex(currentIndex);
        statusDTO.setCurrentDocName(currentDocName);
        statusDTO.setProgress(progress);
        statusDTO.setMessage(message);

        String key = getTaskKey(taskId);
        redisCacheUtil.setCacheObject(key, statusDTO, TASK_STATUS_EXPIRE_HOURS, TimeUnit.HOURS);
    }

    /**
     * 添加已完成的文档
     *
     * @param taskId  任务ID
     * @param docName 文档名称
     */
    public void addCompletedDoc(String taskId, String docName) {
        DocumentParseTaskStatusDTO statusDTO = getTaskStatus(taskId);
        if (statusDTO == null) {
            log.warn("任务状态不存在，taskId: {}", taskId);
            return;
        }

        if (statusDTO.getCompletedDocs() == null) {
            statusDTO.setCompletedDocs(new ArrayList<>());
        }
        statusDTO.getCompletedDocs().add(docName);

        String key = getTaskKey(taskId);
        redisCacheUtil.setCacheObject(key, statusDTO, TASK_STATUS_EXPIRE_HOURS, TimeUnit.HOURS);
    }

    /**
     * 添加失败的文档
     *
     * @param taskId       任务ID
     * @param docName      文档名称
     * @param errorMessage 错误信息
     */
    public void addFailedDoc(String taskId, String docName, String errorMessage) {
        DocumentParseTaskStatusDTO statusDTO = getTaskStatus(taskId);
        if (statusDTO == null) {
            log.warn("任务状态不存在，taskId: {}", taskId);
            return;
        }

        if (statusDTO.getFailedDocs() == null) {
            statusDTO.setFailedDocs(new ArrayList<>());
        }
        statusDTO.getFailedDocs().add(new DocumentParseTaskStatusDTO.FailedDoc(docName, errorMessage));

        String key = getTaskKey(taskId);
        redisCacheUtil.setCacheObject(key, statusDTO, TASK_STATUS_EXPIRE_HOURS, TimeUnit.HOURS);
    }

    /**
     * 标记任务为成功
     *
     * @param taskId 任务ID
     */
    public void markTaskSuccess(String taskId) {
        DocumentParseTaskStatusDTO statusDTO = getTaskStatus(taskId);
        if (statusDTO == null) {
            log.warn("任务状态不存在，taskId: {}", taskId);
            return;
        }

        statusDTO.setStatus(DocumentParseTaskStatus.SUCCESS);
        statusDTO.setMessage("所有文档解析完成");
        statusDTO.setProgress(100);
        statusDTO.setCompleteTime(LocalDateTime.now());

        String key = getTaskKey(taskId);
        redisCacheUtil.setCacheObject(key, statusDTO, TASK_STATUS_EXPIRE_HOURS, TimeUnit.HOURS);
        log.info("文档解析任务完成，taskId: {}", taskId);
    }

    /**
     * 标记任务为失败
     *
     * @param taskId       任务ID
     * @param errorMessage 错误信息
     */
    public void markTaskFailed(String taskId, String errorMessage) {
        DocumentParseTaskStatusDTO statusDTO = getTaskStatus(taskId);
        if (statusDTO == null) {
            log.warn("任务状态不存在，taskId: {}", taskId);
            return;
        }

        statusDTO.setStatus(DocumentParseTaskStatus.FAILED);
        statusDTO.setMessage(errorMessage);
        statusDTO.setCompleteTime(LocalDateTime.now());

        String key = getTaskKey(taskId);
        redisCacheUtil.setCacheObject(key, statusDTO, TASK_STATUS_EXPIRE_HOURS, TimeUnit.HOURS);
        log.error("文档解析任务失败，taskId: {}, error: {}", taskId, errorMessage);
    }

    /**
     * 获取任务状态
     *
     * @param taskId 任务ID
     * @return 任务状态DTO
     */
    public DocumentParseTaskStatusDTO getTaskStatus(String taskId) {
        String key = getTaskKey(taskId);
        return redisCacheUtil.getCacheObject(key);
    }

    /**
     * 获取任务状态key
     *
     * @param taskId 任务ID
     * @return Redis key
     */
    private String getTaskKey(String taskId) {
        return TASK_STATUS_KEY_PREFIX + taskId;
    }
}

